"""
Alpha076-101因子使用示例
======================

演示如何使用Alpha101因子库中的Alpha076-101因子。
包括数据准备、因子计算、结果分析等完整流程。
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import sys
import os
from typing import Dict, List
import warnings

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from alpha101.core import Alpha101Factory, AlphaFactor
from alpha101.config import get_factor_config, list_factors_by_category

# 设置中文字体和忽略警告
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
warnings.filterwarnings('ignore')


def generate_sample_data(n_stocks: int = 100, n_days: int = 252) -> Dict[str, np.ndarray]:
    """
    生成示例股票数据
    
    Parameters:
    -----------
    n_stocks : int
        股票数量
    n_days : int
        交易日数量
        
    Returns:
    --------
    Dict[str, np.ndarray]
        股票数据字典
    """
    print(f"生成示例数据: {n_stocks}只股票, {n_days}个交易日")
    
    np.random.seed(42)
    
    # 生成基础价格数据
    base_price = 100.0
    returns = np.random.normal(0.0005, 0.02, (n_stocks, n_days))
    
    # 添加一些趋势和周期性
    trend = np.linspace(0, 0.1, n_days)
    cycle = 0.05 * np.sin(2 * np.pi * np.arange(n_days) / 50)
    returns += (trend + cycle).reshape(1, -1)
    
    # 生成价格序列
    prices = np.zeros((n_stocks, n_days))
    prices[:, 0] = base_price * (1 + np.random.normal(0, 0.1, n_stocks))
    
    for i in range(1, n_days):
        prices[:, i] = prices[:, i-1] * (1 + returns[:, i])
    
    # 生成OHLC数据
    close = prices
    open_price = close * (1 + np.random.normal(0, 0.005, (n_stocks, n_days)))
    
    # 确保high >= max(open, close), low <= min(open, close)
    high_base = np.maximum(open_price, close)
    low_base = np.minimum(open_price, close)
    
    high = high_base * (1 + np.abs(np.random.normal(0, 0.01, (n_stocks, n_days))))
    low = low_base * (1 - np.abs(np.random.normal(0, 0.01, (n_stocks, n_days))))
    
    # 生成成交量数据（对数正态分布）
    volume = np.random.lognormal(10, 1, (n_stocks, n_days))
    
    # 生成VWAP数据
    vwap = (high + low + close) / 3.0 + np.random.normal(0, 0.001, (n_stocks, n_days))
    
    # 计算收益率
    returns_data = np.zeros((n_stocks, n_days))
    returns_data[:, 1:] = (close[:, 1:] - close[:, :-1]) / close[:, :-1]
    
    # 计算20日平均成交量
    adv20 = np.zeros((n_stocks, n_days))
    for i in range(19, n_days):
        adv20[:, i] = np.mean(volume[:, i-19:i+1], axis=1)
    adv20[:, :19] = np.nan
    
    return {
        'open': open_price,
        'high': high,
        'low': low,
        'close': close,
        'volume': volume,
        'vwap': vwap,
        'returns': returns_data,
        'adv20': adv20
    }


def demonstrate_single_factor():
    """演示单个因子的计算和分析"""
    print("\n" + "="*60)
    print("1. 单个因子计算演示")
    print("="*60)
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # 生成数据
    data = generate_sample_data(n_stocks=50, n_days=100)
    
    # 选择一个因子进行演示
    factor_id = "alpha076"
    
    print(f"计算因子: {factor_id}")
    
    # 获取因子信息
    config = get_factor_config(factor_id)
    print(f"因子描述: {config['description']}")
    print(f"因子类别: {config['category']}")
    print(f"时间窗口: {config['window']}")
    print(f"数据依赖: {config['dependencies']}")
    
    # 计算因子
    factor_values = factory.compute_factor(factor_id, data)
    
    print(f"因子值形状: {factor_values.shape}")
    print(f"有效值数量: {np.sum(np.isfinite(factor_values))}")
    print(f"NaN值数量: {np.sum(np.isnan(factor_values))}")
    
    # 统计信息
    finite_values = factor_values[np.isfinite(factor_values)]
    if len(finite_values) > 0:
        print(f"因子值统计:")
        print(f"  均值: {np.mean(finite_values):.6f}")
        print(f"  标准差: {np.std(finite_values):.6f}")
        print(f"  最小值: {np.min(finite_values):.6f}")
        print(f"  最大值: {np.max(finite_values):.6f}")
        print(f"  中位数: {np.median(finite_values):.6f}")
    
    return factor_values


def demonstrate_batch_computation():
    """演示批量因子计算"""
    print("\n" + "="*60)
    print("2. 批量因子计算演示")
    print("="*60)
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # 生成数据
    data = generate_sample_data(n_stocks=30, n_days=80)
    
    # 选择要计算的因子
    factor_ids = [f"alpha{i:03d}" for i in range(76, 86)]  # 前10个因子
    
    print(f"批量计算因子: {factor_ids}")
    
    # 批量计算
    import time
    start_time = time.time()
    
    results = factory.batch_compute(factor_ids, data)
    
    computation_time = time.time() - start_time
    
    print(f"批量计算完成，耗时: {computation_time:.3f}秒")
    
    # 统计结果
    successful_factors = [fid for fid, result in results.items() if result is not None]
    failed_factors = [fid for fid, result in results.items() if result is None]
    
    print(f"成功计算: {len(successful_factors)} 个因子")
    print(f"失败计算: {len(failed_factors)} 个因子")
    
    if failed_factors:
        print(f"失败的因子: {failed_factors}")
    
    # 分析成功的因子
    if successful_factors:
        print(f"\n成功因子的统计信息:")
        for factor_id in successful_factors[:5]:  # 只显示前5个
            factor_values = results[factor_id]
            finite_values = factor_values[np.isfinite(factor_values)]
            if len(finite_values) > 0:
                print(f"  {factor_id}: 均值={np.mean(finite_values):.4f}, "
                      f"标准差={np.std(finite_values):.4f}, "
                      f"有效值={len(finite_values)}")
    
    return results


def demonstrate_factor_analysis():
    """演示因子分析"""
    print("\n" + "="*60)
    print("3. 因子分析演示")
    print("="*60)
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # 生成数据
    data = generate_sample_data(n_stocks=40, n_days=60)
    
    # 计算几个因子
    factor_ids = ["alpha076", "alpha077", "alpha078", "alpha079", "alpha080"]
    results = {}
    
    print("计算选定因子...")
    for factor_id in factor_ids:
        try:
            factor_values = factory.compute_factor(factor_id, data)
            results[factor_id] = factor_values
            print(f"✓ {factor_id} 计算成功")
        except Exception as e:
            print(f"✗ {factor_id} 计算失败: {e}")
    
    if len(results) < 2:
        print("可用因子不足，跳过相关性分析")
        return
    
    # 计算因子间相关性
    print(f"\n计算 {len(results)} 个因子的相关性...")
    
    # 提取最后一天的因子值
    factor_data = {}
    for factor_id, factor_values in results.items():
        last_day_values = factor_values[:, -1]
        finite_mask = np.isfinite(last_day_values)
        if np.sum(finite_mask) > 10:  # 至少10个有效值
            factor_data[factor_id] = last_day_values[finite_mask]
    
    if len(factor_data) >= 2:
        # 找到共同的有效股票
        min_length = min(len(values) for values in factor_data.values())
        
        # 构建相关性矩阵
        factor_matrix = np.array([values[:min_length] for values in factor_data.values()])
        correlation_matrix = np.corrcoef(factor_matrix)
        
        print("因子相关性矩阵:")
        factor_names = list(factor_data.keys())
        corr_df = pd.DataFrame(correlation_matrix, 
                              index=factor_names, 
                              columns=factor_names)
        print(corr_df.round(3))
        
        # 找出高相关性的因子对
        high_corr_pairs = []
        for i in range(len(factor_names)):
            for j in range(i+1, len(factor_names)):
                corr_val = correlation_matrix[i, j]
                if abs(corr_val) > 0.5:
                    high_corr_pairs.append((factor_names[i], factor_names[j], corr_val))
        
        if high_corr_pairs:
            print(f"\n高相关性因子对 (|相关系数| > 0.5):")
            for factor1, factor2, corr in high_corr_pairs:
                print(f"  {factor1} - {factor2}: {corr:.3f}")
        else:
            print("\n未发现高相关性因子对")


def demonstrate_factor_categories():
    """演示因子分类分析"""
    print("\n" + "="*60)
    print("4. 因子分类分析")
    print("="*60)
    
    # 统计Alpha076-101的因子分类
    factor_ids = [f"alpha{i:03d}" for i in range(76, 102)]
    categories = {}
    
    for factor_id in factor_ids:
        try:
            config = get_factor_config(factor_id)
            category = config['category']
            if category not in categories:
                categories[category] = []
            categories[category].append(factor_id)
        except:
            continue
    
    print("Alpha076-101因子分类统计:")
    for category, factors in categories.items():
        print(f"  {category}: {len(factors)} 个因子")
        print(f"    {factors[:5]}{'...' if len(factors) > 5 else ''}")
    
    # 分析每个类别的特点
    print(f"\n各类别因子特点:")
    for category, factors in categories.items():
        windows = []
        for factor_id in factors:
            try:
                config = get_factor_config(factor_id)
                windows.append(config['window'])
            except:
                continue
        
        if windows:
            avg_window = np.mean(windows)
            print(f"  {category}: 平均时间窗口 {avg_window:.1f} 天")


def demonstrate_performance_comparison():
    """演示性能比较"""
    print("\n" + "="*60)
    print("5. 性能比较演示")
    print("="*60)
    
    # 创建因子工厂
    factory = Alpha101Factory(use_numba=False)
    
    # 生成不同规模的数据进行性能测试
    test_configs = [
        (20, 50, "小规模"),
        (50, 100, "中规模"),
        (100, 200, "大规模")
    ]
    
    # 选择几个代表性因子
    test_factors = ["alpha076", "alpha080", "alpha085", "alpha090", "alpha095"]
    
    print("性能测试结果:")
    print(f"{'数据规模':<10} {'因子':<10} {'计算时间(秒)':<12} {'内存使用':<10}")
    print("-" * 50)
    
    for n_stocks, n_days, scale_name in test_configs:
        data = generate_sample_data(n_stocks, n_days)
        
        for factor_id in test_factors:
            try:
                import time
                start_time = time.time()
                
                factor_values = factory.compute_factor(factor_id, data)
                
                computation_time = time.time() - start_time
                memory_usage = factor_values.nbytes / 1024 / 1024  # MB
                
                print(f"{scale_name:<10} {factor_id:<10} {computation_time:<12.3f} {memory_usage:<10.2f}MB")
                
            except Exception as e:
                print(f"{scale_name:<10} {factor_id:<10} {'ERROR':<12} {str(e)[:20]}")


def main():
    """主函数"""
    print("Alpha076-101因子使用示例")
    print("=" * 60)
    
    try:
        # 1. 单个因子计算演示
        demonstrate_single_factor()
        
        # 2. 批量因子计算演示
        demonstrate_batch_computation()
        
        # 3. 因子分析演示
        demonstrate_factor_analysis()
        
        # 4. 因子分类分析
        demonstrate_factor_categories()
        
        # 5. 性能比较演示
        demonstrate_performance_comparison()
        
        print("\n" + "="*60)
        print("演示完成！")
        print("="*60)
        
        print("\n使用建议:")
        print("1. 在实际使用中，建议开启Numba加速以提高性能")
        print("2. 对于大规模数据，可以考虑分批计算")
        print("3. 注意处理NaN值和无穷大值")
        print("4. 建议对因子值进行标准化和去极值处理")
        print("5. 可以根据因子相关性进行因子选择")
        
    except Exception as e:
        print(f"演示过程中发生错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()